import torch
import torch.nn as nn
import numpy as np
import math
import matplotlib.pyplot as plt


class Embedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.embed = nn.Embedding(num_embeddings=self.vocab_size, embedding_dim=d_model)

    def forward(self, x):
        return self.embed(x) * math.sqrt(self.d_model)


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=60):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(size=(max_len, d_model))
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = 1 / (10000 ** (torch.arange(0, d_model, 2).float() / d_model))
        position_value = position * div_term
        pe[:, 0::2] = torch.sin(position_value)
        pe[:, 1::2] = torch.cos(position_value)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

    def forward(self, embed):
        result = embed + self.pe[:, : embed.shape[1]]
        return self.dropout(result)


def plot_position():
    my_position = PositionalEncoding(20, dropout=0.1, max_len=100)
    embed = torch.zeros(1, 100, 20)
    y = my_position(embed)
    plt.figure(figsize=(20, 15))
    plt.plot(np.arange(100), y[0, :, 4:8].detach().numpy())
    plt.legend([f"dim {p}" for p in [4, 5, 6, 7]])
    plt.show()
